from .base import PoseEstimator, center_ellipse
import numpy as np
from pose import initial_pose


class EPnP(PoseEstimator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __name__(self):
        return "epnp"

    def forward(self, C_3d, C_2d,*args, **ignore_kwargs):
        x_cnt = []
        X_cnt = []
        for c_3d, c_2d in zip(C_3d, C_2d):
            # 求解相机椭圆中心：
            x_cnt.append(center_ellipse(c_2d))
            X_cnt.append(center_ellipse(c_3d))
        # 求解相机位姿
        X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
        x_cnt = np.array(x_cnt)
        return self.estimator(X_cnt, x_cnt)

    def estimator(self, database, measurements):
        """
        K: 相机内参
        database: 三维点库
        measurements: 观测到的二维点
        """
        assert len(database) == len(measurements)
        assert measurements.shape[1] == 2
        assert database.shape[1] == 3
        # 求解相机位姿
        P_ba = initial_pose(self.K, measurements, database)
        if P_ba[0] is None:
            return False, None, None
        t = P_ba[1]  # 提取平移部分
        R = P_ba[0]  # 提取旋转部分
        return True, R, t
